import os.path
import time
import torch
import torch.nn as nn
import torch.optim as optim
from torchsummary import summary
from data import load_cifar10_data
from models.model import create_model


# 计算准确率
def accuracy(y_hat, y):
    return float((y_hat.argmax(axis=1) == y).sum())


# 验证模型
def evaluate_accuracy(net, val_loader, device):
    net.eval()  # 设置为评估模式
    # 正确预测的数量，总预测的数量
    eval_acc_sum, n = 0.0, 0
    with torch.no_grad():
        for X, y in val_loader:
            X = X.to(device)
            y = y.to(device)
            y_hat = net(X)
            eval_acc_sum += accuracy(y_hat, y)
            n += y.numel()
    return eval_acc_sum / n 


# 训练模型
def train_one_epoch(epoch, model, train_loader, val_loader, criterion, optimizer, device):
    model.to(device)
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

        if batch_idx % 100 == 0:
            val_acc = evaluate_accuracy(model, val_loader, device)
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}\tVal Acc: {:.4f}'
                .format(epoch, batch_idx * len(data), len(train_loader.dataset),
                        100. * batch_idx / len(train_loader), loss.item(), val_acc))
            

# 验证模型
def test(model, test_loader, criterion, device):
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    test_acc = correct / len(test_loader.dataset)
    print('Validation: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)'.format(
        test_loss, correct, len(test_loader.dataset), 100. * test_acc), end='')
    return (test_acc, test_loss, )


# 整合训练步骤
def train(args, device, in_channels, num_classes):

    # 实例化模型
    model = create_model(args.model, in_channels, num_classes, use_timm=False)
    if (args.print_model):
        summary(model, (in_channels, model.image_size, model.image_size), device='cpu')

    if (args.checkpoint is not None):
        if (os.path.exists(args.checkpoint)):
            model.load_state_dict(torch.load(args.checkpoint))
            print('Model checkpoint loaded: {}'.format(args.checkpoint))

    if (args.train):
        # 加载数据集
        train_loader, val_loader, test_loader, _ = load_cifar10_data(args.batch_size, model.image_size)

        # 定义损失函数和优化器
        criterion = nn.CrossEntropyLoss()
        # optimizer = optim.SGD(model.parameters(), lr=lr)
        optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
        # optimizer = optim.Adam(model.parameters(), lr=lr)
        # optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=wd)

        scheduler = None
        if (args.sched is not None):
            scheduler = optim.lr_scheduler.StepLR(optimizer, args.lr_period, args.lr_decay)

        # 运行训练和验证
        for epoch in range(args.epochs):
            start_time = time.time()
            train_one_epoch(epoch, model, train_loader, val_loader, criterion, optimizer, device)
            if (args.eval):
                res = test(model, test_loader, criterion, device)
            if (scheduler is not None):
                scheduler.step()
                print(', lr={:.4f}'.format(scheduler.get_last_lr()[0]), end='')
            print(', Elasped time: {:.4f}\n'.format(time.time() - start_time))

            if args.output is not None:
                if not os.path.exists(args.output):
                    os.mkdir(args.output)
                filepath = os.path.join(args.output, 'checkpoint_{}.pt'.format(epoch))
                torch.save(model.state_dict(), filepath)
                print('Model checkpoint saved: {}'.format(filepath))